import abc
from dataclasses import dataclass

from .action_scheduling import ActionScheduling


@dataclass(kw_only=True)
class LearningRateSchedulingHparams:
    learning_rate: float | None = None
    warmup_duration: float | int | None = None
    decay_halflife: float | int | None = None
    decay_scale: float | None = None


class LearningRateSchedulingTarget(abc.ABC):
    @property
    @abc.abstractmethod
    def learning_rate(self) -> float:
        pass

    @learning_rate.setter
    @abc.abstractmethod
    def learning_rate(self, value: float):
        pass


class LearningRateScheduling(ActionScheduling):
    def __init__(self, *,
        progress_metric_name: str = 'step',
        action_interval: int | float | None = 1,
        target: LearningRateSchedulingTarget | None = None,
        hparams: LearningRateSchedulingHparams,
    ):
        super().__init__(
            progress_metric_name=progress_metric_name,
            action_interval=action_interval, action_mode='pre')
        self._target: LearningRateSchedulingTarget | None = target
        self._hparams: LearningRateSchedulingHparams = hparams

    # action

    def _action(self, progress: int | float):
        if self._target is None:
            self._target = self.dock.get_plugin(LearningRateSchedulingTarget)
        learning_rate_scale = self._compute_learning_rate_scale(self._hparams, progress)
        self._target.learning_rate = self._hparams.learning_rate * learning_rate_scale

    # learning rate scale

    @staticmethod
    def _compute_learning_rate_scale(hparams: LearningRateSchedulingHparams, progress: float) -> float:
        warmup_duration = hparams.warmup_duration if hparams.warmup_duration is not None else 0
        decay_halflife = hparams.decay_halflife
        decay_scale = hparams.decay_scale if hparams.decay_scale is not None else 0
        if progress < warmup_duration:
            return progress / warmup_duration
        if decay_halflife is not None:
            return decay_scale + (1 - decay_scale) * \
                2 ** (-(progress - warmup_duration) / decay_halflife)
        return 1.0
